# -*- coding: utf-8 -*-

from __future__ import unicode_literals
from __future__ import print_function
from __future__ import division
from __future__ import absolute_import

from datetime import datetime
from operator import add


###############################################################################
# common functions                                                            #
###############################################################################

def parse_date(s):
    return datetime.strptime(s, "%Y-%m-%d")


def year_from_date(s):
    if not s or not isinstance(s, basestring):
        return -1
    try:
        return datetime.now().year - parse_date(s).year
    except ValueError:
        return -1


def is_within_one_year(s):
    y = year_from_date(s)
    return 0 <= y <= 1


def is_list_contains_keywords(values, keywords):
    for k in keywords:
        if k in values:
            return True
    return False


def is_list_contains_different_values(values):
    return len(set(values)) > 1


###############################################################################
# lambdas                                                                     #
###############################################################################

def str_add(x, y):
    return "{0};{1}".format(x, y)


def sum_pair(left, right):
    return left[0] + right[0], left[1] + right[1]


###############################################################################
# rdd utilities                                                               #
###############################################################################

def count_by_key(pair_rdd, num_partitions=None):
    """
    for key value rdd, count the number of values for each key
    optionally distinct the rdd if number of partitions is specified
    :param pair_rdd: [key => value]
    :param num_partitions:
    :return: [key => count]
    """
    if num_partitions is not None:
        pair_rdd = pair_rdd.distinct(num_partitions)

    return pair_rdd.mapValues(lambda x: 1).reduceByKey(add)


def reduce_avg(rdd, num_partitions):
    """
    for key value rdd, calculate the average of each key
    :param rdd: [key => value]
    :param num_partitions: the final number of partitions
    :return:
    """
    return rdd.map(lambda x: (x[0], (x[1], 1))).reduceByKey(sum_pair, num_partitions).mapValues(
        lambda x: x[0] / x[1])


def group_by_key(rdd):
    """
    aggregate by key and group relevant values with a set

    :param rdd: [key => value]
    :return: [key => set(values)]
    """

    def update_and_return(x, y):
        x.add(y)
        return x

    return rdd.aggregateByKey(set(), update_and_return, lambda x, y: x.union(y))


def count_value_by_key(rdd, num_partitions):
    """
    for key value rdd, count number of different values by key
    :param rdd: [key => value]
    :param num_partitions: number of partitions when groupByKey
    :return: [key => [(value1, count1), (value2, count2), ...] where count1 >= count2 >= ...
    """
    return rdd.map(lambda x: ((x[0], x[1]), 1)).reduceByKey(add).map(
        lambda x: (x[0][0], (x[0][1], x[1]))).groupByKey(num_partitions).mapValues(
        lambda x: [i for i in sorted(x, key=lambda t: t[1], reverse=True)])


def get_top_n_values(rdd, n):
    """
    for key value rdd, for each key, get the top n values and concat them as a string
    :param rdd: [key => [(value1, count1), (value2, count2), ...] where count1 >= count2 >= ...
    :param n: integer
    :return: [key => "value_1;value_2;value_3;...;value_n"]
    """
    return rdd.mapValues(lambda x: reduce(str_add, [i[0] for i in x[:n]]))


###############################################################################
# kg utilities                                                                #
###############################################################################

def is_of_type(kg_item, type_str):
    return kg_item["@type"][0] == type_str


def have_type(kg_item, type_str):
    return type_str in kg_item["@type"]


def extract_entity_of_type(entity, _type, predicate=have_type, persist=False):
    """
    extract entities of specific type from a [kgid => entity] rdd

    :param entity: a collection of [kgid => entity]
    :param _type: the wanted type signature of an entity
    :param predicate: how to decide if an entity is of some type
    :param persist: if persist is required for the rdd
    :return: rdd of the required type
    """
    res = entity.filter(lambda x: predicate(x[1], _type))
    if persist:
        res = res.persist()
    return res


def extract_link_of_type(link, _type, predicate=have_type, persist=False):
    """
    extract links of specific type from a [link] rdd

    :param link: a collection of [link]
    :param _type: the wanted type signature of a link
    :param predicate: how to decide if an entity is of some type
    :param persist: if persist is required for the rdd
    :return: rdd of the required type
    """
    res = link.filter(lambda x: predicate(x, _type))
    if persist:
        res = res.persist()
    return res


def merge_entity_attributes(entity_rdd, entity_attr_rdd, attribute_name=None):
    """
    merge attributes of entity_attr_rrd to entity_rdd
    :param entity_rdd: [id => entity]
    :param entity_attr_rdd: [id => attribute] or [id => json]
    :param attribute_name: attribute_name to set
    :return: [id => entity] with updated attributes
    """

    def update_json(item):
        if item[1] is not None:
            item[0].update(item[1])
        return item[0]

    def update_attr(item):
        if item[1] is not None:
            item[0][attribute_name] = item[1]
        return item[0]

    if attribute_name:
        update_func = update_attr
    else:
        update_func = update_json

    return entity_rdd.leftOuterJoin(entity_attr_rdd).mapValues(update_func)
